from confluent_kafka.admin import AdminClient
from confluent_kafka import Consumer, TopicPartition


def get_lag_info(broker, group_id, topic):
    """
    获取指定消费者组在指定主题上的LAG信息
    """
    try:
        admin_client = AdminClient({'bootstrap.servers': broker})
        consumer = Consumer({'bootstrap.servers': broker, 'group.id': group_id})
        # 获取主题分区信息
        topic_metadata = consumer.list_topics(topic).topics[topic]
        partitions = topic_metadata.partitions

        topic_partitions = [TopicPartition(topic, p) for p in partitions.keys()]
        lag_info = []

        try:
            # 尝试获取消费者组提交的偏移量
            committed_offsets = consumer.committed(topic_partitions, timeout=5.0)
        except Exception:
            committed_offsets = []

        offset_dict = {}
        for tp in committed_offsets:
            if tp.offset >= 0:  # 有效偏移量
                offset_dict[tp.partition] = tp.offset
            else:
                offset_dict[tp.partition] = None

        for tp in topic_partitions:
            try:
                # 获取lag End offset(分区最新偏移量)
                low, high = consumer.get_watermark_offsets(tp, timeout=5.0)
                # 获取消费者组提交的当前偏移量
                current_offset = offset_dict.get(tp.partition)

                lag = high - current_offset if current_offset is not None else 'Unknown'
                lag_info.append({
                    'partition': tp.partition,
                    'current_offset': current_offset,
                    'lag_end_offset': high,
                    'lag': lag
                })
            except Exception as e:
                # 详细记录错误信息
                lag_info.append({
                    'partition': tp.partition,
                    'current_offset': None,
                    'lag_end_offset': None,
                    'lag': f'Error: {str(e)}',
                    'error_details': f'Failed to get watermark offsets for partition {tp.partition}'
                })
    except Exception as e:
        print(f"连接Kafka或获取元数据时出错: {e}")
        return []
    finally:
        try:
            consumer.close()
        except:
            pass

    return lag_info
